from model.resnet101 import resnet101
import torch.nn as nn
import torch

class Resnet101_FCN8s(nn.Module):
    def __init__(self,n_class):
        super(Resnet101_FCN8s, self).__init__()
        self.encode = resnet101(pretrained=True)
        self.trans_p5 = nn.Conv2d(2048,n_class,1)
        self.trans_p4 = nn.Conv2d(1024,n_class,1)
        self.trans_p3 = nn.Conv2d(512,n_class,1)

        self.smooth_conv1 = nn.Conv2d(n_class,n_class,3,padding=1)
        self.smooth_conv2 = nn.Conv2d(n_class,n_class,3,padding=1)


        self.up2time = nn.ConvTranspose2d(n_class,n_class,2,stride=2,bias=False)
        self.up4time = nn.ConvTranspose2d(n_class,n_class,2,stride=2,bias=False)
        self.up8time = nn.ConvTranspose2d(n_class,n_class,8,stride=8,bias=False)

    def forward(self, x):
        p1,p2, p3, p4, p5 = self.encode(x)
        f5 = self.trans_p5(p5)
        up2time = self.up2time(f5)

        h = self.trans_p4(p4)
        h +=up2time
        h = self.smooth_conv1(h)

        up4time = self.up4time(h)
        h = self.trans_p3(p3)
        h+=up4time
        h = self.smooth_conv2(h)

        final_scores = self.up8time(h)
        return final_scores


if __name__ == '__main__':
    model = Resnet101_FCN8s(12)
    model.eval()
    x = torch.randn(1,3,352,480)
    y = model(x)
    print(y.shape)